from tool import check_tpu_type


# TPU工厂函数，根据不同的TPU传回不同实例
def TPUFactory(algorithm, tpu_type=None):

    if tpu_type is None:
        # TPU的类型检测，采用检测其供应商ID和产品ID的方式
        tpu_type_list = check_tpu_type()
        if len(tpu_type_list) == 0:
            print("未检测到加速棒")
            exit(0)
        elif len(tpu_type_list) >= 1:
            tpu_type = tpu_type_list[0]

    if tpu_type == "songke":
        if algorithm == "yolov5":
            from songke_1000.predict_yolov5 import PredictYolov5
            return PredictYolov5()
        elif algorithm == "yolov8":
            from songke_1000.predict_yolov8 import PredictYolov8
            return PredictYolov8()
        elif algorithm == "mcts":
            from songke_1000.predict_mcts import PredictMcts
            return PredictMcts()
        elif algorithm == "obb":
            from songke_1000.predict_obb import PredictObb
            return PredictObb()

    elif tpu_type == "kneron":
        if algorithm == "yolov5":
            from kneron_520.predict_yolov5 import PredictYolov5
            return PredictYolov5()
        elif algorithm == "yolov8":
            from kneron_520.predict_yolov8 import PredictYolov8
            return PredictYolov8()
        elif algorithm == "mcts":
            from kneron_520.predict_mcts import PredictMcts
            return PredictMcts()
    else:
        print("该加速棒类型不存在")
